import faiss
import numpy as np
import matplotlib.pyplot as plt

d = 1024  #维数
n_data = 2000
np.random.seed(0)
data = []
mu = 3
sigma = 0.1
for i in range(n_data):
    data.append(np.random.normal(mu, sigma, d))
data = np.array(data).astype('float32')

query = []
n_query = 10
np.random.seed(0)
for i in range(n_query):
    query.append(np.random.normal(mu, sigma, d))
query = np.array(query).astype('float32')
# IndexFlatL2类型遍历计算查询向量与被查询向量的L2精确距离，
# 不需要训练操作（大部分index类型都需要train操作）。
index = faiss.IndexFlatL2(d)
# print(index.is_trained)
# print(index.ntotal)
index.add(data)
# print(index.ntotal)
# ----------------------------------------------------------------------------
'''精确索引'''
# 返回结果数
k = 3
query_self = data[:5]
dis,ind = index.search(query_self,k)
# 升序返回每个查询向量的距离
print(dis)
# 升序返回每个查询向量的k个相似结果
print(ind)

'''倒排表快速索引'''
nlist = 50  # 将数据库向量分割为多少了维诺空间
k = 10
quantizer = faiss.IndexFlatL2(d)  # 量化器
index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
       # METRIC_L2计算L2距离, 或faiss.METRIC_INNER_PRODUCT计算内积
assert not index.is_trained   #倒排表索引类型需要训练
index.train(data)  # 训练数据集应该与数据库数据集同分布
assert index.is_trained

index.add(data)
index.nprobe = 50  # 选择n个维诺空间进行索引,
dis, ind = index.search(query, k)
print(dis)
print(ind)

'''乘积量化索引'''
'''在上述两种索引方式中，在index中都保存了完整的数据库向量，
在数据量非常大的时候会占用太多内存，甚至超出内存限制。
在faiss中，当数据量非常大的时候，一般采用乘积量化方法保存原始向量的有损压缩形式,
故而查询阶段返回的结果也是近似的。'''

nlist = 50
m = 8                             # 列方向划分个数，必须能被d整除
k = 10
quantizer = faiss.IndexFlatL2(d)
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 4)
                                    # 4 表示每个子向量被编码为 4 bits
index.train(data)
index.add(data)
index.nprobe = 50
dis, ind = index.search(query_self, k)  # 查询自身
print(dis)
print(ind)
dis, ind = index.search(query, k)  # 真实查询
print(dis)
print(ind)